from flask import Blueprint, jsonify, request
import akshare as ak
import pandas as pd
from datetime import datetime

trading_days = Blueprint('trading_days', __name__)

@trading_days.route('/trading_days/count', methods=['GET'])
def get_trading_days_count():
    """
    获取从指定时间点开始的A股交易日数量
    
    参数:
        start_date (str): 开始日期，格式为 'YYYY-MM-DD' (必需)
        end_date (str): 结束日期，格式为 'YYYY-MM-DD' (可选，默认为今天)
    
    返回:
        JSON: 包含交易日数量
    """
    try:
        # 获取请求参数
        start_date = request.args.get('start_date')
        end_date = request.args.get('end_date')
        
        # 验证必需参数
        if not start_date:
            return jsonify({
                'success': False,
                'error': 'start_date参数是必需的',
                'message': '请提供开始日期，格式为YYYY-MM-DD'
            }), 400
        
        # 验证日期格式
        try:
            datetime.strptime(start_date, '%Y-%m-%d')
            if end_date:
                datetime.strptime(end_date, '%Y-%m-%d')
        except ValueError:
            return jsonify({
                'success': False,
                'error': '日期格式错误',
                'message': '日期格式应为YYYY-MM-DD'
            }), 400
        
        # 获取历史交易日数据
        trading_dates_df = ak.tool_trade_date_hist_sina()
        trading_dates_df['trade_date'] = pd.to_datetime(trading_dates_df['trade_date'])
        
        # 转换输入日期
        start_dt = pd.to_datetime(start_date)
        if end_date:
            end_dt = pd.to_datetime(end_date)
        else:
            end_dt = datetime.now()
        
        # 过滤指定日期范围内的交易日
        filtered_dates = trading_dates_df[
            (trading_dates_df['trade_date'] >= start_dt) & 
            (trading_dates_df['trade_date'] <= end_dt)
        ]
        
        # 计算交易日数量
        trading_days_count = len(filtered_dates)
        
        return jsonify({
            'success': True,
            'start_date': start_date,
            'end_date': end_dt.strftime('%Y-%m-%d'),
            'trading_days_count': trading_days_count
        })
        
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e),
            'message': '获取交易日数据失败'
        }), 500 